#coding:utf-8
'''
心智识别
Author: lxw
Date: 2023.12.06
'''
import numpy as np
import tensorflow as tf
import os
import torch
import face_recognition
from tensorflow.keras import layers
import requests
from io import BytesIO
import uuid
from alignment import get_crop_img
from upload_oss import upload_file

model = ''
model_w = './models/class_test_W_1.h5'
model_c = './models/class_test_C_1.h5'
model_l = './models/class_test_L.h5'
model_g = './models/class_test_G_1.h5'
def init():
    global model  # 声明要修改的是全局变量 model
    # 忽略tf的warning
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    # 禁用TensorFlow的警告信息
    # tf.get_logger().setLevel('ERROR')
    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if torch.cuda.is_available():
        print('use GPU')
    else:
        print('use CPU')

    # 单指数代码
    # 下面的代码用来创建神经网络
    model = tf.keras.Sequential()
    model.add(tf.keras.Input(shape=(128,)))

    model.add(layers.Dense(768))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('relu'))
    # model.add(layers.Dropout(0.1))

    model.add(layers.Dense(512))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('relu'))
    # model.add(layers.Dropout(0.1))

    model.add(layers.Dense(256))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('relu'))
    # model.add(layers.Dropout(0.1))

    model.add(layers.Dense(128))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('relu'))
    # model.add(layers.Dropout(0.1))

    model.add(layers.Dense(63))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('relu'))
    # model.add(layers.Dropout(0.1))

    model.add(layers.Dense(61))
    model.add(layers.BatchNormalization())
    model.add(layers.Activation('softmax'))

    




def predict(img_path):
    result = {}
    try:
        #初始化
        init()
        response = requests.get(img_path)  # 发起网络请求获取图片数据
        # 创建保存图片的文件夹
        save_folder = 'face_crop'
        os.makedirs(save_folder, exist_ok=True)

        # 发起网络请求获取图片数据
        response = requests.get(img_path)
        img_data = response.content

        # 生成保存图片的文件路径和文件名
        # 生成唯一的文件名
        # unique_filename = str(uuid.uuid4())
        filename = str(uuid.uuid4())+'.jpg'
        save_path = os.path.join(save_folder, filename)

        # 将图片数据写入文件
        with open(save_path, 'wb') as file:
            file.write(img_data)

        # 获取图片的相对地址
        relative_path = os.path.relpath(save_path)

        #进行图片裁剪优化
        ret = get_crop_img(relative_path)
        if ret["code"] != 0:
            result["code"] = 101
            result["message"] = ret["msg"]
            return result


        # print(relative_path)
        img = face_recognition.load_image_file(ret["data"])
        output = face_recognition.face_encodings(img, model='large')
    except:
        result["code"] = 101
        result["message"] = "图片格式有误或图片损坏无法读取"
        return result
    else:
        # print('output',output)
        face_num = len(output)
        # print('共有X张人脸',face_num)
        if face_num == 0:
            result["code"] = 200
            result["message"] = "图片中未检测到人脸"
            return result
        elif face_num > 1:
            result["code"] = 201
            result["message"] = "图片中包含多个人脸"
            return result
        elif face_num == 1:
            output = output[0] * 10
            output = output.reshape(1, 128)
            # 多指数预测
            # 心智指数预测
            # pred = model.predict(output) * 100
            # W_pred = int(np.round(pred[0, 0]))
            # C_pred = int(np.round(pred[0, 1]))
            # L_pred = int(np.round(pred[0, 2]))
            # G_pred = int(np.round(pred[0, 3]))
            # 单指数预测
            # model = tf.keras.models.load_model(model_w) 
            model.load_weights(model_w)
            W_pred_origin = model.predict(output)[0]
            W_pred = 0
            for idx in range(len(W_pred_origin)):
                W_pred += W_pred_origin[idx] * (idx + 20)

            model.load_weights(model_c)
            #model = tf.keras.models.load_model(model_c)
            C_pred_origin = model.predict(output)[0]
            C_pred = 0
            for idx in range(len(C_pred_origin)):
                C_pred += C_pred_origin[idx] * (idx + 20)

            model.load_weights(model_l)
            #model = tf.keras.models.load_model(model_l)
            L_pred_origin = model.predict(output)[0]
            L_pred = 0
            for idx in range(len(L_pred_origin)):
                L_pred += L_pred_origin[idx] * (idx + 20)

            model.load_weights(model_g)
            #model = tf.keras.models.load_model(model_g)
            G_pred_origin = model.predict(output)[0]
            G_pred = 0
            for idx in range(len(G_pred_origin)):
                G_pred += G_pred_origin[idx] * (idx + 20)
            # result
            result["code"] = 500
            result["message"] = "图片识别成功"
            result["data"] = {}
            result["data"]["W"] = int(np.round(W_pred))
            result["data"]["C"] = int(np.round(C_pred))
            result["data"]["L"] = int(np.round(L_pred))
            result["data"]["G"] = int(np.round(G_pred))

            # 将图片上传到阿里云
            crop_img =  upload_file(ret["data"])
            if crop_img == '':
                result["code"] = 101
                result["message"] = "图片存储失败"
                return result
            result["data"]["crop_img"] = crop_img
            # print(result)
            return result

  
# 校验
# print("\ncheck face model")
# print(predict('http://langfang-img.oss-cn-beijing.aliyuncs.com/3898f302-4be3-448f-94e1-bf47196afc39.jpg'))